import os
import cv2
from PIL import Image, ImageDraw
import math
import numpy as np
try:
    from .cvio import cvio
except:
    from cvio import cvio


class ColorMAP:

    def __init__(self, shapes):
        self.color_map = self.gen_color_map(shapes)

    def gen_color_map(self, shapes):
        color_map = {}
        for shape in shapes:
            label = shape['label']
            if label in color_map:
                continue
            color_mask = tuple(np.random.randint(0, 255, (3), dtype=np.uint8).tolist())
            color_map[label] = color_mask
        return color_map

    def __call__(self, shape):
        return self.color_map[shape['label']]


def gray2rgb(mask, color):
    # mask is a binary img array (0, 1)
    mask = np.concatenate([(mask * c)[:, :, None]
                           for c in color], axis=2).astype(np.uint8)
    return mask


def get_rect_loc(shape, loc='upper'):
    polygon = shape['points']
    polygon = np.array(polygon, dtype=np.int32)
    assert loc in ('upper', 'lower')
    xmin = polygon[:, 0].min()
    ymin = polygon[:, 1].min()
    xmax = polygon[:, 0].max()
    ymax = polygon[:, 1].max()
    bbox = ((xmin, ymin), (xmax, ymax))
    if loc == 'upper':
        label_bbox = ((xmin, ymin), (xmax, ymin + 15))
    else:
        label_bbox = ((xmin, ymax - 15), (xmax, ymax))
    return bbox, label_bbox

def custom_draw_shapes(img, shapes, mode='contour', alpha=0.8, label_text='upper'):
    assert mode in ('rectangle', 'contour', 'mask')
    assert alpha >= 0 and alpha <= 1
    if isinstance(img, Image.Image):
        img = np.array(img, dtype=np.uint8)
    color_map = ColorMAP(shapes)
    for i, shape in enumerate(shapes):
        color = color_map(shape)
        label = shape['label']
        rect_loc, text_loc = get_rect_loc(shape, label_text)
   
        shape_type = shape['shape_type']

        if mode in ('contour', 'mask') and alpha > 0:
            if shape_type != 'polygon':
                continue            
            (x1, y1), (x2, y2) = rect_loc
            roi = img[y1: y2, x1:x2]
            h, w, c = roi.shape
            if x1 >= x2 or y1 >= y2 or h <= 1 or w <= 1:
                print('Exception', shape)
                continue
            mask = np.zeros((h, w), dtype=np.uint8)
            mask = Image.fromarray(mask)
            draw = ImageDraw.Draw(mask)
            
            xy = [tuple([int(p[0] - x1), int(p[1] - y1)]) for p in shape['points']]
            draw.polygon(xy=xy, outline=1, fill=1)
            mask = np.array(mask, dtype=np.uint8)
            for i, c in enumerate(color):
                c = 255 - c
                if mode == 'contour':
                    kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(5, 5))
                    dilated = cv2.dilate(mask, kernel).astype(np.uint8)
                    mask = dilated - mask         
                inds = np.where(mask > 0)           
                img[y1: y2, x1:x2, i][inds] = np.array(roi[..., i] * (1 - alpha) + mask * c * alpha, dtype=np.uint8)[inds]

        cv2.rectangle(img, rect_loc[0], rect_loc[1], color, 2, cv2.LINE_AA)
        cv2.rectangle(img, text_loc[0], text_loc[1],
                      (255, 255, 255), -1, cv2.LINE_AA)  # filled
        label_loc = (text_loc[0][0], text_loc[0][1] + 10)        
        # text_color = tuple([255 - c for c in color])
        text_color = (0,0,0)
        cv2.putText(img, label, label_loc, 0, 0.35,
                    text_color, thickness=1, lineType=cv2.LINE_AA)     

    return Image.fromarray(img)

def plot_mask(src, dst, label_text='upper', alpha=0.5, mode='contour'):
    if not os.path.exists(dst):
        os.makedirs(dst)
    img_ann_list = cvio.load_img_ann_list(src)
    alpha = 1 - alpha
    for i, (imgsrc, annsrc) in enumerate(img_ann_list,1):
        print('[%d/%d] %s' % (i, len(img_ann_list), imgsrc))
        try:
            img = Image.open(imgsrc)
            ann_info = cvio.load_ann(annsrc)
            shapes = ann_info['shapes']
            img = custom_draw_shapes(img, shapes, mode=mode, alpha=alpha, label_text=label_text)
            img.save(os.path.join(dst, os.path.basename(imgsrc)))
        except Exception as e:
            print(e)

if __name__ == '__main__':
    src = r'G:\data\datasets\drink\daiding_101\base_test\jt\cascade_mask_x101_fpn_dcn_fl\plt'
    dst = r'G:\data\datasets\drink\daiding_101\base_test\jt\cascade_mask_x101_fpn_dcn_fl\plt'
    plot_mask(src, dst, label_text='upper', mode='mask', alpha=0.35)
